
// "Copyright [2021] <Copyright Shizhen Zhao>"

#include <algorithm>
#include <cstdlib>

#include "inet/node/inet/MmuManager.h"

namespace inet {

MmuManager::MmuManager(int num_ports) : port_count_(num_ports), hdrm_bytes_(num_ports, 0),
    ingress_bytes_(num_ports, 0), egress_bytes_(num_ports, 0), paused_(num_ports, 0),
    kmax_(num_ports, 0), kmin_(num_ports, 0), headroom_(num_ports, 0) {
    shared_used_bytes_ = 0;
}

void MmuManager::ConfigMmu(const std::vector<int>& bandwidth_gbps,
                          const std::vector<int>& link_delay_ns) {
    
    // Configure mmu size
    buffer_size_ = 12 * 1024 * 1024;  // Switch is shadow buffered, with only 12MB memory.

    // Configure headroom
    total_hdrm_ = 0;
    for (int port = 0; port < port_count_; ++port) {
        headroom_[port] = bandwidth_gbps[port] * link_delay_ns[port] / 8 * 3 + 1024 * 5;
        total_hdrm_ += headroom_[port];
    }

    // Configure reserve
    reserve_ = 4 * 1024;  // Each port has 4KB reserved memory.
    total_rsrv_ = reserve_ * port_count_;

    // Configure pause and resume related variables
    // TODO(shizhenzhao): changes pfc_a_shift_ based on port bandwidth.
    pfc_a_shift_ = 3;
    resume_offset_ = 3 * 1024;  // Resume offset is about 2 MTU in DCQCN paper.

    // Configure ECN related variables
    pmax_ = 0.2;
    for (int port = 0; port < port_count_; ++port) {
        kmax_[port] = std::max(bandwidth_gbps[port] * 64, 6400);
        kmin_[port] = std::max(bandwidth_gbps[port] * 16, 1600);
    }
}

bool MmuManager::CheckIngressAdmission(int in_port, int packet_size_byte) {
    if (packet_size_byte + hdrm_bytes_[in_port] > headroom_[in_port] &&
        packet_size_byte + GetSharedUsed(in_port) > GetPfcThreshold()) {
        return false;
    }
    return true;
}

void MmuManager::UpdateIngressAdmission(int in_port, int packet_size_byte) {
    int new_bytes = ingress_bytes_[in_port] + packet_size_byte;
    if (new_bytes <= reserve_) {
        ingress_bytes_[in_port] += packet_size_byte;
    } else {
        int thresh = GetPfcThreshold();
        if (new_bytes - reserve_ > thresh){
            hdrm_bytes_[in_port] += packet_size_byte;
        }else {
            ingress_bytes_[in_port] += packet_size_byte;
            shared_used_bytes_ += std::min(packet_size_byte, new_bytes - reserve_);
        }
    }
}

void MmuManager::UpdateEgressAdmission(int out_port, int packet_size_byte) {
    egress_bytes_[out_port] += packet_size_byte;
}

void MmuManager::RemoveFromIngressAdmission(int in_port, int packet_size_byte) {
    int from_hdrm = std::min(hdrm_bytes_[in_port], packet_size_byte);
    int from_shared = std::min(packet_size_byte - from_hdrm,
        ingress_bytes_[in_port] > reserve_ ? ingress_bytes_[in_port] - reserve_ : 0);
    hdrm_bytes_[in_port] -= from_hdrm;
    ingress_bytes_[in_port] -= packet_size_byte - from_hdrm;
    shared_used_bytes_ -= from_shared;
}

void MmuManager::RemoveFromEgressAdmission(int out_port, int packet_size_byte) {
    egress_bytes_[out_port] -= packet_size_byte;
}

bool MmuManager::ShouldSendECN(int out_port) {
    if (egress_bytes_[out_port] > kmax_[out_port])
        return true;
    if (egress_bytes_[out_port] > kmin_[out_port]) {
        double p = pmax_ * double(egress_bytes_[out_port] - kmin_[out_port]) / (kmax_[out_port] - kmin_[out_port]);
        double random_variable = double(rand()) / RAND_MAX;
        if (random_variable < p)
            return true;
    }
    return false;
}

bool MmuManager::CheckShouldPause(int in_port) {
    return !paused_[in_port] && (hdrm_bytes_[in_port] > 0 ||
        GetSharedUsed(in_port) >= GetPfcThreshold());
}

bool MmuManager::CheckShouldResume(int in_port){
    if (!paused_[in_port])
        return false;
    int shared_used = GetSharedUsed(in_port);
    return hdrm_bytes_[in_port] == 0 && (shared_used == 0 ||
        shared_used + resume_offset_ <= GetPfcThreshold());
}

void MmuManager::SetPause(int in_port) {
    paused_[in_port] = true;
}

void MmuManager::SetResume(int in_port) {
    paused_[in_port] = false;
}

int MmuManager::GetPfcThreshold() {
    return (buffer_size_ - total_hdrm_ - total_rsrv_ - shared_used_bytes_) >> pfc_a_shift_;
}

int MmuManager::GetSharedUsed(int in_port) {
    int used = ingress_bytes_[in_port];
    return used > reserve_ ? used - reserve_ : 0;
}

}  // namespace inet
